"""
根据物体坐标点和像素坐标点估算物体姿态，并建立坐标系

1. 加载相机内参矩阵、畸变系数
2. 加载图片
3. 查找每个图片的角点
4. 查找角点亚像素
5. 计算对象姿态solvePnpRansac
6. 投影3D点到图像平面
7. 在图片上坐标系并显示图片
"""
import numpy as np
import cv2
import glob


def draw(img, corners, imgpts):
    corner = tuple(corners[0].ravel())
    func = lambda x: [int(i) for i in x]
    img = cv2.line(img, func(corner), func(tuple(imgpts[0].ravel())), (255,0,0), 5)
    img = cv2.line(img, func(corner), func(tuple(imgpts[1].ravel())), (0,255,0), 5)
    img = cv2.line(img, func(corner), func(tuple(imgpts[2].ravel())), (0,0,255), 5)
    return img

if __name__ == '__main__':
    with open("cameraMatrix.txt", 'r') as f:
        data = f.read().splitlines()

    mtx = np.array(eval(data[0]), dtype=np.float32).reshape(3, 3)
    dist = np.array(eval(data[1]), dtype=np.float32)
    print(mtx)
    print(dist)


    objp = np.zeros((6 * 9, 3), np.float32)
    objp[:, :2] = np.mgrid[0:6, 0:9].T.reshape(-1, 2)


    criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.001)

    """
    [[ 3.  0.  0.]
     [ 0.  3.  0.]
     [ 0.  0. -3.]]
    """
    axis = np.float32([[2, 0, 0], [0, 2, 0], [0, 0, -2]]).reshape(-1, 3)

    img = cv2.imread("40.jpg")
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    ret, corners = cv2.findChessboardCorners(gray, (6, 9), None)

    if ret:
        corners2 = cv2.cornerSubPix(gray, corners, (11, 11), (-1, -1), criteria)

        # 根据对象点和角点列表，查找旋转向量和平移向量
        # retval, rvecs, tvecs, inliers = cv2.solvePnPRansac(objp, corners2, mtx, dist)
        retval, rvecs, tvecs = cv2.solvePnP(objp, corners2, mtx, dist)

        # 将3D点投影到图像平面
        # axis 为所需要的float类型的3D点列表
        # 输出图像点和雅克比矩阵
        imagePoints, jacobian = cv2.projectPoints(axis, rvecs, tvecs, mtx, dist)

        img = draw(img, corners2, imagePoints)
        cv2.imshow('img', img)
        cv2.waitKey(0)
        # cv2.imwrite("axis_result.jpg", img)

    cv2.destroyAllWindows()